import numpy as np
import struct
import os
import faiss
import multiprocessing


# to get the .vecs
# np.set_printoptions(threshold=np.inf)  # display all the content when print the numpy array


def ivecs_read(fname):
    a = np.fromfile(fname, dtype='int32')
    d = a[0]
    return a.reshape(-1, d + 1)[:, 1:].copy(), d


def fvecs_read(fname):
    data, d = ivecs_read(fname)
    return data.view('float32').astype(np.float32), d


def bvecs_read(fname):
    a = np.fromfile(fname, dtype='uint8')
    d = a[:4].view('uint8')[0]
    return a.reshape(-1, d + 4)[:, 4:].copy(), d


# store in format of vecs
def fvecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))  # *dimension就是int, dimension就是list
        f.write(struct.pack('f' * len(x), *x))

    f.close()


def ivecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))
        f.write(struct.pack('i' * len(x), *x))

    f.close()


def bvecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))
        f.write(struct.pack('B' * len(x), *x))

    f.close()


def ip_gnd(base, query, k):
    base_dim = base.shape[1]
    index = faiss.IndexFlatIP(base_dim)
    index.add(base)
    gnd_distance, gnd_idx = index.search(query, k)
    return gnd_idx, gnd_distance


def l2_gnd(base, query, k):
    base_dim = base.shape[1]
    index = faiss.IndexFlatL2(base_dim)
    index.add(base)
    gnd_distance, gnd_idx = index.search(query, k)
    return gnd_idx, gnd_distance


def normalization(vectors):
    vecs_module = np.linalg.norm(vectors, axis=1)
    vecs_module = vecs_module[:, np.newaxis]
    vectors = vectors / vecs_module
    return vectors


def cosine_gnd(base, query, k):
    base_dim = base.shape[1]
    base_norm = normalization(base)
    query_norm = normalization(query)
    index = faiss.IndexFlatL2(base_dim)
    index.add(base_norm)
    gnd_distance, gnd_idx = index.search(query_norm, k)
    return gnd_idx, gnd_distance


def norm_gnd(base, query, k):
    norm_base = np.linalg.norm(base, axis=1)
    topk_idx = np.argsort(-norm_base)[:k]
    topk_dist = norm_base[topk_idx]
    return np.array([topk_idx]), np.array([topk_dist])


def delete_dir_if_exist(dire):
    if os.path.isdir(dire):
        command = 'rm -rf %s' % dire
        print(command)
        os.system(command)


def get_bias_data(old_basic_dir, new_basic_dir, ds_name, topk):
    old_base_fname = '%s/%s/%s_base.fvecs' % (old_basic_dir, ds_name, ds_name)
    old_query_fname = '%s/%s/%s_query.fvecs' % (old_basic_dir, ds_name, ds_name)
    base, dim = fvecs_read(old_base_fname)
    query, dim = fvecs_read(old_query_fname)

    delete_dir_if_exist("%s/%s_bias" % (new_basic_dir, ds_name))
    os.mkdir("%s/%s_bias" % (new_basic_dir, ds_name))

    ip_gnd_idx, _ = ip_gnd(base, query, topk)
    gnd_idx_fname = '%s/%s_bias/ip_gnd_idx.ivecs' % (new_basic_dir, ds_name)
    ivecs_write(gnd_idx_fname, ip_gnd_idx)
    print("write %s_bias ip groundtruth" % ds_name)

    l2_gnd_idx, _ = l2_gnd(base, query, len(base))
    gnd_idx_fname = '%s/%s_bias/l2_gnd_idx.ivecs' % (new_basic_dir, ds_name)
    ivecs_write(gnd_idx_fname, l2_gnd_idx)
    print("write %s_bias l2 groundtruth" % ds_name)
    del ip_gnd_idx, l2_gnd_idx

    cosine_gnd_idx, _ = cosine_gnd(base, query, len(base))
    gnd_idx_fname = '%s/%s_bias/cosine_gnd_idx.ivecs' % (new_basic_dir, ds_name)
    ivecs_write(gnd_idx_fname, cosine_gnd_idx)
    print("write %s_bias cosine groundtruth" % ds_name)

    norm_gnd_idx, _ = norm_gnd(base, query, len(base))
    gnd_idx_fname = '%s/%s_bias/norm_gnd_idx.ivecs' % (new_basic_dir, ds_name)
    ivecs_write(gnd_idx_fname, norm_gnd_idx)
    print("write %s_bias norm groundtruth" % ds_name)
    del cosine_gnd_idx, norm_gnd_idx


if __name__ == '__main__':
    dataset_l = ['audio', 'netflix']
    # dataset_l = ['word2vec', 'yahoomusic']
    # dataset_l = ['text-to-image', 'tiny5m',
    #              'word2vec', 'yahoomusic']
    topk = 10
    old_basic_dir = '/home/zhengbian/Dataset/inner_product'
    new_basic_dir = '/home/zhengbian/mips-graph/data'
    for dataset_name in dataset_l:
        get_bias_data(old_basic_dir, new_basic_dir, dataset_name, topk)
